import React, { useState, useEffect } from 'react';
import axios from 'axios';

interface Dataset {
  id: string;
  name: string;
  description?: string;
  indexing_technique?: string;
  retrieval_model?: any;
}

interface RetrievalResult {
  content: string;
  score: number;
  source: {
    dataset_id: string;
    document_id: string;
    segment_id: string;
    document_name?: string;
    data_source_type?: string;
  };
}

const RetrievalTest: React.FC = () => {
  const [datasets, setDatasets] = useState<Dataset[]>([]);
  const [selectedDataset, setSelectedDataset] = useState<string>('');
  const [query, setQuery] = useState<string>('');
  const [topK, setTopK] = useState<number>(3);
  const [results, setResults] = useState<RetrievalResult[]>([]);
  const [loading, setLoading] = useState<boolean>(false);
  const [error, setError] = useState<string | null>(null);

  useEffect(() => {
    fetchDatasets();
  }, []);

  const fetchDatasets = async () => {
    try {
      const response = await axios.get<{ datasets: Dataset[] }>('/api/datasets'); // Assuming API is proxied
      setDatasets(response.data.datasets);
      if (response.data.datasets.length > 0) {
        setSelectedDataset(response.data.datasets[0].id);
      }
    } catch (err) {
      console.error('Error fetching datasets:', err);
      setError('Failed to fetch datasets.');
    }
  };

  const handleSearch = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!selectedDataset || !query) {
      setError('Please select a dataset and enter a query.');
      return;
    }
    setLoading(true);
    setError(null);
    setResults([]);
    try {
      const response = await axios.get<{ results: RetrievalResult[] }>(`/api/datasets/${selectedDataset}/query`, {
        params: { q: query, top_k: topK },
      });
      setResults(response.data.results);
    } catch (err: any) {
      console.error('Error performing retrieval:', err);
      setError(err.response?.data?.error || 'Failed to perform retrieval.');
    } finally {
      setLoading(false);
    }
  };

  return (
    <div style={{ padding: '20px', fontFamily: 'Arial, sans-serif' }}>
      <h2>Test Knowledge Base Retrieval</h2>
      <form onSubmit={handleSearch}>
        <div style={{ marginBottom: '15px' }}>
          <label htmlFor="dataset-select" style={{ marginRight: '10px' }}>Select Dataset:</label>
          <select 
            id="dataset-select" 
            value={selectedDataset} 
            onChange={(e) => setSelectedDataset(e.target.value)}
            style={{ padding: '8px', minWidth: '200px' }}
            disabled={datasets.length === 0}
          >
            {datasets.map((ds) => (
              <option key={ds.id} value={ds.id}>{ds.name} ({ds.id.substring(0,8)})</option>
            ))}
          </select>
        </div>
        <div style={{ marginBottom: '15px' }}>
          <label htmlFor="query-input" style={{ marginRight: '10px' }}>Query:</label>
          <input 
            type="text" 
            id="query-input" 
            value={query} 
            onChange={(e) => setQuery(e.target.value)} 
            placeholder="Enter your query"
            style={{ padding: '8px', width: 'calc(100% - 100px)', minWidth: '300px' }}
          />
        </div>
        <div style={{ marginBottom: '15px' }}>
          <label htmlFor="topk-input" style={{ marginRight: '10px' }}>Top K:</label>
          <input 
            type="number" 
            id="topk-input" 
            value={topK} 
            onChange={(e) => setTopK(parseInt(e.target.value, 10))} 
            min="1"
            style={{ padding: '8px', width: '60px' }}
          />
        </div>
        <button type="submit" disabled={loading || !selectedDataset || !query} style={{ padding: '10px 15px', cursor: 'pointer' }}>
          {loading ? 'Searching...' : 'Search'}
        </button>
      </form>

      {error && <p style={{ color: 'red', marginTop: '15px' }}>Error: {error}</p>}

      {results.length > 0 && (
        <div style={{ marginTop: '20px' }}>
          <h3>Results:</h3>
          {results.map((result, index) => (
            <div key={index} style={{ border: '1px solid #ccc', padding: '10px', marginBottom: '10px', borderRadius: '5px' }}>
              <p><strong>Score:</strong> {result.score !== undefined && result.score !== null ? result.score.toFixed(4) : 'N/A'}</p>
              <p style={{ whiteSpace: 'pre-wrap' }}><strong>Content:</strong> {result.content}</p>
              <p style={{ fontSize: '0.9em', color: '#555' }}>
                Source: Dataset ID: {result.source.dataset_id.substring(0,8)}, Document: {result.source.document_name || result.source.document_id.substring(0,8)}, Segment ID: {result.source.segment_id.substring(0,8)}
              </p>
            </div>
          ))}
        </div>
      )}
    </div>
  );
};

export default RetrievalTest;

